import os
import cv2
import numpy as np
from rknnlite.api import RKNNLite
import traceback
import sys

# 常量配置
IMG_SIZE = 224
RKNN_MODEL = 'food.rknn'
IMG_FOLDER = 'img'
LABEL_FILE = 'foodlabel.txt'  # 使用foodlabel.txt作为标签文件


def load_class_labels(label_file):
    """从标签文件加载类别名称"""
    try:
        # 检查文件是否存在
        if not os.path.exists(label_file):
            print(f"警告: 标签文件 '{label_file}' 不存在")
            return []

        # 使用UTF-8编码读取文件
        with open(label_file, "r", encoding="utf-8") as f:
            class_names = [line.strip() for line in f.readlines()]

        print(f"从 '{label_file}' 加载了 {len(class_names)} 个类别")
        return class_names

    except Exception as e:
        print(f"加载类别标签时出错: {str(e)}")
        return []


def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


def preprocess(img):
    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    return np.expand_dims(img, axis=0)


def postprocess(output, class_names):
    probs = softmax(output[0]).squeeze()
    class_id = np.argmax(probs).item()
    confidence = probs[class_id].item()

    # 确保类别索引在有效范围内
    if class_id < len(class_names):
        class_name = class_names[class_id]
    else:
        class_name = f"unknown_{class_id}"

    return class_name, confidence


def save_result_image(img_path, class_name, confidence):
    """保存带有识别结果的图片"""
    img = cv2.imread(img_path)
    font = cv2.FONT_HERSHEY_SIMPLEX
    text = f"{class_name}: {confidence:.4f}"
    position = (10, 50)
    font_scale = 1
    color = (0, 255, 0)  # Green color
    thickness = 2

    cv2.putText(img, text, position, font, font_scale, color, thickness)
    output_img_path = os.path.join(os.path.dirname(img_path), f"result_{os.path.basename(img_path)}")
    cv2.imwrite(output_img_path, img)
    print(f"结果图片已保存到: {output_img_path}")


def rknn_main():
    try:
        print(f"当前工作目录: {os.getcwd()}")
        print(f"Python sys.path: {sys.path}")

        # 加载类别标签
        class_names = load_class_labels(LABEL_FILE)
        if not class_names:
            # 如果无法加载标签文件，使用默认值
            print("使用默认类别列表")
            class_names = ["moldy", "new"]

        rknn = RKNNLite()

        # 加载RKNN模型
        print('加载RKNN模型...')
        ret = rknn.load_rknn(RKNN_MODEL)
        if ret != 0:
            print('加载RKNN模型失败!')
            return ["error"]

        # 初始化运行时
        print('初始化运行时环境...')
        ret = rknn.init_runtime()
        if ret != 0:
            print('初始化运行时失败!')
            return ["error"]

        # 获取img文件夹中的第一张图片
        img_files = os.listdir(IMG_FOLDER)
        if not img_files:
            print("错误: 在img文件夹中没有找到图片")
            return ["error"]

        img_path = os.path.join(IMG_FOLDER, img_files[0])
        print(f"处理图片: {img_path}")
        img = cv2.imread(img_path)

        if img is None:
            print(f"错误: 无法读取图片 {img_path}")
            return ["error"]

        # 预处理
        input_data = preprocess(img)

        # 执行推理
        print('执行推理...')
        outputs = rknn.inference(inputs=[input_data])

        # 后处理 - 使用从文件加载的类别信息
        class_name, confidence = postprocess(outputs, class_names)

        print(f"识别结果: {class_name}, 置信度: {confidence:.4f}")

        # 保存结果图片
        save_result_image(img_path, class_name, confidence)

        # 释放资源
        rknn.release()
        return [class_name]

    except Exception as e:
        print(f"运行过程中出错: {str(e)}")
        traceback.print_exc()
        return ["error"]


# 测试代码
if __name__ == '__main__':
    print("作为独立Python脚本运行")
    results = rknn_main()
    print("最终结果:", results)

    # 确保输出为UTF-8编码
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except:
        pass  # 对于旧版Python可能不支持reconfigure

    # 打印结果时显式指定编码
    print("最终结果 (UTF-8):", results)



